import torch
import numpy as np
from matplotlib import pyplot as plt
from sklearn.datasets import load_boston
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split

boston = load_boston()
data = boston.data
target = boston.target.reshape(-1, 1)

scaler_data = MinMaxScaler()
scaler_target = MinMaxScaler()
data = scaler_data.fit_transform(data)
target = scaler_target.fit_transform(target)

c = 7
x = []
y = []
for i in range(len(data) - c):
    x_data = data[i:i + c]
    y_data = target[i + c]
    x.append(x_data)
    y.append(y_data)

x = torch.Tensor(x)
y = torch.Tensor(y)

train_x, test_x, train_y, test_y = train_test_split(x, y, shuffle=False)
print(train_x.shape)

from torch import nn


class HousePricePredictionCNN(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Sequential(
            nn.Conv1d(in_channels=train_x.shape[2], out_channels=64, kernel_size=2),
            nn.ReLU())
        self.conv2 = nn.Sequential(
            nn.Conv1d(in_channels=64, out_channels=64, kernel_size=2),
            nn.ReLU())
        self.flat = nn.Flatten()
        self.fc = nn.Sequential(
            nn.Linear(in_features=320, out_features=512),
            nn.ReLU(),
            nn.Linear(in_features=512, out_features=1024),
            nn.ReLU(),
            nn.Linear(in_features=1024, out_features=1)
        )

    def forward(self, x):
        x = x.permute(0, 2, 1)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.flat(x)
        x = self.fc(x)
        return x


# 创建模型实例
model = HousePricePredictionCNN()

# 编译模型
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练模型
epochs = 1000

for epoch in range(epochs):
    model.train()
    optimizer.zero_grad()
    y_pred = model(train_x)
    train_y = train_y.view(-1, 1)
    loss = criterion(y_pred, train_y)
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        print(f'Epoch {epoch + 1}, Loss: {loss.item()}')

model.eval()
with torch.no_grad():
    y_pred = model(test_x).squeeze()
    plt.plot(y_pred.numpy(), label='Predicted', c='r')
    plt.plot(test_y.numpy(), label='Actual', c='g')
    plt.legend()
    plt.show()
